from deeprobust.graph.defense import GCN
from deeprobust.graph.global_attack import MetaApprox, Metattack
from deeprobust.graph.utils import *
from load_data import load


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# device = torch.device("cpu")

adj, features, labels, idx_train, idx_val, idx_test, idx_unlabeled = load(name='citeseer')

adj, features, labels = preprocess(adj, features, labels, preprocess_adj=False)
# Setup Surrogate Model
surrogate = GCN(nfeat=features.shape[1], nclass=labels.max().item() + 1, nhid=16,
                dropout=0.5, with_relu=False, with_bias=True, weight_decay=5e-4, device=str(device)).to(device)
surrogate.fit(features, adj, labels, idx_train)


def main(name, nums):
    model = Metattack(model=surrogate, nnodes=adj.shape[0], feature_shape=features.shape,
            attack_structure=True, attack_features=True, device=str(device), lambda_=1).to(device)
    model.attack(features, adj, labels, idx_train, idx_unlabeled, nums, ll_constraint=False)
    print('=== testing GCN on original(clean) graph ===')
    # test(adj)
    adj_table = './outputs/mettack/' + name + '/random/' + str(nums) + '_modified_adj.pt'
    features_table = './outputs/mettack/' + name + '/random/' + str(nums) + '_modified_features.pt'
    torch.save(model.modified_adj, adj_table)
    torch.save(model.modified_features, features_table)
    print('第%d保存成功' % nums)


if __name__ == '__main__':
    data = 'cora'
    num = [0.02, 0.05, 0.1, 0.3, 0.5, 0.8]
    for i in num:
        perturbations = int(i * (adj.sum() // 2))
        main(data, perturbations)
        # print(perturbations)
